#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <geometry_msgs/PoseStamped.h>
#include <std_msgs/String.h>
#include <cv_bridge/cv_bridge.h>
#include <opencv2/opencv.hpp>

#include <iostream>
#include <vector>
#include <string>
#include <stdio.h>
#include <chrono>

using namespace std;
using namespace cv;

//----------------function declaration----------------
void process(Mat& frame);
int detectColor(Mat& frame);
void publish(ros::Publisher& pub);
void print();

//----------------global variables--------------------
struct Target
{
    int    color = 0;      // 1: red, 2: green, 3: blue, 4: yellow
    int    shape = 0;      // 1: square, 2: circle, 3: triangle
    double offset_x = 0.0; // up: 0~100 down: -100~0
    double offset_y = 0.0; // left: 0~100 right: -100~0
};

vector<Target> targets;

//----------------ROS Callback function for image topic----------------
void imageCallback(const sensor_msgs::ImageConstPtr& msg)
{
    try
    {
        // Convert the ROS image message to an OpenCV Mat using cv_bridge
        cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
        Mat frame = cv_ptr->image;

        if (frame.empty())
        {
            ROS_ERROR("Received empty frame");
            return;
        }

        // Process the frame
        process(frame);

        // Show the processed frame (optional for debugging)
        imshow("Frame", frame);
        waitKey(1);
        // print();
        

    }
    catch (cv_bridge::Exception& e)
    {
        ROS_ERROR("cv_bridge exception: %s", e.what());
    }
}

//----------------Main function----------------
int main(int argc, char **argv)
{
    ros::init(argc, argv, "vision_node");
    ros::NodeHandle nh;
    ros::Rate rate(30);

    ros::Publisher vision_pub = nh.advertise<geometry_msgs::PoseStamped>("/ele_vision", 10);
    ros::Subscriber image_sub = nh.subscribe<sensor_msgs::Image>(
        // "/iris/camera/image_raw", 1, imageCallback
        "/usb_cam/image_raw", 1, imageCallback
    );
    ROS_INFO("Started vision node");
    // ros::spin();
    while (ros::ok())
    {
        // Publish target information
        // if(targets.size() > 0)
        // {
        //     publish(vision_pub);
        // }
        publish(vision_pub);
        print();
        ros::spinOnce();

        rate.sleep();
    }

    return 0;
}

//----------------Process the frame----------------
void process(Mat& frame)
{
    Mat gray, thresh;
    vector<vector<Point>> contours;
    vector<Vec4i> hierarchy;

    cvtColor(frame, gray, COLOR_BGR2GRAY);
    Canny(gray, thresh, 69, 200);
    findContours(thresh, contours, hierarchy, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);

    imshow("Thresh", thresh);

    for (size_t i = 0; i < contours.size(); i++)
    {
        double area, bboxArea, ratio;
        int cX, cY;
        RotatedRect minRect;
        Target target;
        Point2f rect_points[4];
        Moments M;
        Mat framez;

        area = contourArea(contours[i]);
        if (area < 3000)
            continue;

        minRect = minAreaRect(contours[i]);
        bboxArea = minRect.size.area();
        ratio = area / bboxArea;

        // Calculate center
        M = moments(contours[i]);
        if (M.m00 != 0)
        {
            cX = int(M.m10 / M.m00);
            cY = int(M.m01 / M.m00);
        }
        else
        {
            cX = frame.cols / 2;
            cY = frame.rows / 2;
        }

        // Store target
        if (ratio > 0.82)
            target.shape = 1; // square
        else if (ratio < 0.82 && ratio > 0.70)
            target.shape = 2; // circle
        else if (ratio < 0.55 && ratio > 0.45)
            target.shape = 3; // triangle
        else
            target.shape = 0;

        if (target.shape)
        {
            framez = Mat::zeros(frame.size(), frame.type());
            Mat maskROI = Mat::zeros(frame.size(), CV_8UC1);
            vector<vector<Point>> currentContour = {contours[i]};
            fillPoly(maskROI, currentContour, Scalar(255));
            bitwise_and(frame, frame, framez, maskROI);
            target.color = detectColor(framez);
        }

        target.offset_x = (frame.rows / 2.0 - cY) / frame.rows;
        target.offset_y = (frame.cols / 2.0 - cX) / frame.cols;

        // Draw rectangle
        minRect.points(rect_points);
        for (int j = 0; j < 4; j++)
            line(frame, rect_points[j], rect_points[(j + 1) % 4], Scalar(0, 255, 0), 2);
        circle(frame, Point(cX, cY), 5, Scalar(255, 0, 0), -1);

        targets.push_back(target);
    }
}

//----------------Publish target information----------------
void publish(ros::Publisher& pub)
{
    for (auto target : targets)
    {
        if (!target.shape || !target.color)
            continue;

        geometry_msgs::PoseStamped pose_msg;
        pose_msg.header.stamp = ros::Time::now();
        pose_msg.header.frame_id = "camera_coordinate";
        pose_msg.pose.position.x = target.color;
        pose_msg.pose.position.y = target.shape;
        pose_msg.pose.position.z = target.offset_x;
        pose_msg.pose.orientation.w = target.offset_y;

        pub.publish(pose_msg);

        ROS_INFO("Published: [Color:%d, Shape:%d, Offset_X:%f, Offset_Y:%f]",
                 target.color, target.shape, target.offset_x, target.offset_y);
    }
    targets.clear();
}

//----------------Detect color----------------
int detectColor(Mat& frame)
{
    Mat hsv, mask_red1, mask_red2, mask_red, mask_yellow, mask_green, mask_blue;
    cvtColor(frame, hsv, COLOR_BGR2HSV);
    Scalar lower_red1(0, 60, 0), upper_red1(22, 255, 255);
    Scalar lower_red2(163, 60, 77), upper_red2(180, 255, 255);
    Scalar lower_yellow(25, 98, 94), upper_yellow(42, 255, 255);
    Scalar lower_green(49, 86, 86), upper_green(73, 255, 255);
    Scalar lower_blue(100, 43, 46), upper_blue(124, 255, 255);

    inRange(hsv, lower_red1, upper_red1, mask_red1);
    inRange(hsv, lower_red2, upper_red2, mask_red2);
    bitwise_or(mask_red1, mask_red2, mask_red);
    inRange(hsv, lower_yellow, upper_yellow, mask_yellow);
    inRange(hsv, lower_green, upper_green, mask_green);
    inRange(hsv, lower_blue, upper_blue, mask_blue);

    imshow("masksss", frame);
    imshow("mask_red", mask_red);
    Mat gray_frame;
    cvtColor(frame, gray_frame, COLOR_BGR2GRAY);
    int totalNonZero = countNonZero(gray_frame);
    int redNonZero = countNonZero(mask_red);
    int yellowNonZero = countNonZero(mask_yellow);
    int greenNonZero = countNonZero(mask_green);
    int blueNonZero = countNonZero(mask_blue);

    double redRatio = static_cast<double>(redNonZero) / totalNonZero;
    double yellowRatio = static_cast<double>(yellowNonZero) / totalNonZero;
    double greenRatio = static_cast<double>(greenNonZero) / totalNonZero;
    double blueRatio = static_cast<double>(blueNonZero) / totalNonZero;

    if (redRatio > 0.8)
        return 1;
    if (yellowRatio > 0.8)
        return 4;
    if (greenRatio > 0.8)
        return 2;
    if (blueRatio > 0.8)
        return 3;

    return 0;
}

//--------------print global varities----------------
void print()
{
    for(auto target : targets)
    {
        cout << "color: "    << target.color    << endl;
        cout << "shape: "    << target.shape    << endl;
        cout << "offset_x: " << target.offset_x << endl;
        cout << "offset_y: " << target.offset_y << endl;
    }
    targets.clear();
}